from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_core.tools import tool
import re
from src.nl2sql.prompts import extrac_prompt


def extract_configuration_tables(schema):
    """
    根据表结构提取外连接表信息
    :param schema:
    :return:
    """
    matchs = re.findall(extrac_prompt, schema)
    return matchs


def get_tables_schema(db, tables_name, llm):
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    tools = toolkit.get_tools()
    get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
    schema_list = list()
    if isinstance(tables_name, list):
        for table_name in tables_name:
            table_schema = get_schema_tool.invoke(table_name)
            schema_list.append(table_schema)
        return schema_list
    table_schema = get_schema_tool.invoke(tables_name)
    return table_schema